#include <iostream>
#include <opencv2/opencv.hpp>
#include <io.h>  // 遍历文件夹文件
#include <opencv2/core/utils/logger.hpp>

#include "trt_run.h"

int main() {
    //LOG_LEVEL_ERROR - 只输出错误日志 / LOG_LEVEL_SILENT - 不在输出日志
    cv::utils::logging::setLogLevel(cv::utils::logging::LOG_LEVEL_ERROR);

    // 图像全路径单斜杠需要转义，或者需要R"()"这样
    // 检测路径
    std::string images_path = R"(E:\new_code\github\dataset\coco128\images\train2017)";
    // 模型路径
    std::string engine_file = R"(G:\Github\yolov8_20231114\ultralytics\my_test\weights\yolov8_seg\yolov8s-seg_int8.engine)";
    // 框和iou置信度
    const float conf_threshold = 0.3f;
    // 检测类别
    const std::vector<std::string> class_names = {"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train",
                                                  "truck", "boat", "traffic light", "fire hydrant", "stop sign",
                                                  "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep",
                                                  "cow", "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella",
                                                  "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard",
                                                  "sports ball", "kite", "baseball bat", "baseball glove", "skateboard",
                                                  "surfboard", "tennis racket", "bottle", "wine glass", "cup", "fork",
                                                  "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange",
                                                  "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair",
                                                  "couch", "potted plant", "bed", "dining table", "toilet", "tv",
                                                  "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave",
                                                  "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase",
                                                  "scissors", "teddy bear", "hair drier", "toothbrush"};
    // 是否画框
    const bool draw_box = true;

    // 存储文件信息的结构体
    struct _finddata_t file_info{};
    std::string current_path = images_path + "/*.jpg";  // 查找特定后缀的文件，支持通配符
    // 第一次查找，查找成功返回longlong型的句柄，失败返回-1
    long long handle = _findfirst(current_path.c_str(), &file_info);

    // 调用构造函数，初始化模型
    YoloDetect yolo_detect(engine_file, conf_threshold, class_names, draw_box);
    // 存储画框图
    cv::Mat img;
    cv::Mat draw_img;
    // 存储输出结果
    std::vector<Detection> result;

    double start, end, t;
    do {
//        // 单张图像路径
        std::string img_path = images_path + "/" + file_info.name;
//
        // Mat是创建一个数组
        // IMREAD_COLOR/1-默认彩色  IMREAD_GRAYSCALE/0-灰度显示
        // IMREAD_UNCHANGED/-1-四通道png
        img = cv::imread(img_path, 1);
        draw_img = img;
//        std::cout << draw_img.type() << std::endl;
//        std::cout<<draw_img.size<<std::endl;  // 2048 x 2448
//        std::cout<<draw_img.channels()<<std::endl;  // 3
//        std::cout<<draw_img.dims<<std::endl;  // 2

        start = clock();
//        for (int i = 0; i < 100; i++) {
//            // 侦测
//            result = yolo_detect.detect(draw_img);
//        }
        // 侦测
        result = yolo_detect.detect(draw_img);

        end = clock();
        // 计算毫秒ms
        t = (end - start) / CLOCKS_PER_SEC * 1000;
        std::cout << "detect_time: " << t << std::endl;

        std::cout << "out_boxes: " << result.size() << std::endl;
//        for (const Detection &detection: result) {
//            int clas_index = detection.clas_id;
//            const std::string& label = class_names[clas_index];
//            float conf = (float) std::round(detection.conf * 100) / 100;
//            int x = detection.box.x;
//            int y = detection.box.y;
//            int w = detection.box.width;
//            int h = detection.box.height;
//
//            std::cout << label << ": " << conf << " box(x,y,w,h): " << x << ", " << y << ", " << w << ", " << h
//                      << std::endl;
//        }

        std::cout << "============================" << std::endl;

//        cv::namedWindow("show_img", cv::WINDOW_NORMAL);
//        cv::imshow("show_img", draw_img);
//        cv::waitKey(0);

    } while (!_findnext(handle, &file_info));  // 查找下一个
    _findclose(handle);  // 关闭

    cv::destroyAllWindows();  // 销毁窗口
    system("pause");
}